from tensorflow.keras.layers import Input
import tensorflow as tf
 
 
def window_partition(x,window_size):
    _,H,W,C = x.shape.as_list()
    # print(H,W,C)
    # input('zz')
    x = tf.reshape(x,shape=[-1,H//window_size,window_size,
                            W//window_size,window_size,C])
    
    # -> B,nH,nW,w,w,C
    x = tf.transpose(x,[0,1,3,2,4,5])
    windows = tf.reshape(x,shape=[-1,window_size,window_size,C])
    return windows
 
def window_reverse(windows,window_size,H,W,C):
    # print(f'in window_reverse, {windows.shape}')
    # print(windows.shape,window_size,H,W,C)
    # input('zz')
    x = tf.reshape(windows,shape=[-1,H//window_size,W//window_size,
                                window_size,window_size,C])
    
    x = tf.transpose(x,[0,1,3,2,4,5])
    x = tf.reshape(x,shape=[-1,H,W,C])
    return x
 
def drop_path(inputs,drop_prob,is_training=None):
    if (not is_training) or (drop_prob==0.):
        return inputs
    keep_prob = 1.0 - drop_prob
 
    random_tensor = keep_prob
    shape = (tf.shape(inputs)[0],) + (1,)*(len(tf.shape(inputs))-1)
    random_tensor +=(tf.random.uniform(shape,dtype=inputs.dtype))
    binary_tensor = tf.floor(random_tensor)
    output = tf.math.divide(inputs,keep_prob) * binary_tensor
    return output
 
 
# if __name__ == '__main__':
#     inputs = Input(shape=[56,56,96],batch_size=2)
#     windows = window_partition(inputs,7)
#     print(windows.shape)
#     x = window_reverse(windows,7,56,56,96)
#     print(x.shape)
 